package org.example;



import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;

import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;

import java.util.Arrays;

public class Main {
    public static void main(String[] args) throws IOException {

        //创建一个tcp服务端
        int port = 1080;
        Selector selector = Selector.open();
        ServerSocketChannel serverSocketChannel = ServerSocketChannel.open();
        serverSocketChannel.bind(new InetSocketAddress(port));
        serverSocketChannel.configureBlocking(false);
        serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);

        System.out.println("tcp服务端启动成功,端口："+port);


        while (true) {
            selector.select();
            for (SelectionKey key : selector.selectedKeys()) {
                try{
                    if (key.isAcceptable()) {
                        //服务端接收客户端请求
                        SocketChannel socketChannel = serverSocketChannel.accept();
                        socketChannel.configureBlocking(false);
                        socketChannel.register(selector, SelectionKey.OP_READ,1);
                    } else if(key.isConnectable()){
                        //本地连接
                        SocketChannel socketChannel = (SocketChannel) key.channel();
                        Object[] attachment = (Object[])key.attachment();
                        SocketChannel clientChannel = (SocketChannel)attachment[0];
                        byte[] replyData = (byte[])attachment[1];
                        try {
                            // 完成连接[这里会出现超时异常]
                            socketChannel.finishConnect();
                            //两个channel交换数据
                            socketChannel.register(selector,SelectionKey.OP_READ,clientChannel);
                            //重新注册事件
                            clientChannel.register(selector, SelectionKey.OP_READ,socketChannel);
                            //告诉客户端连接成功
                            clientChannel.write(ByteBuffer.wrap(replyData));

                        } catch (Exception ex) {
                            ex.printStackTrace();
                            socketChannel.close();
                            clientChannel.close();
                        }
                    }else if (key.isReadable()) {
                        SocketChannel socketChannel = (SocketChannel) key.channel();
                        Object attachment = key.attachment();
                        ByteBuffer buffer = ByteBuffer.allocate(1024);
                        try{
                            int bytesRead = socketChannel.read(buffer);
                            if (bytesRead == -1) {
                                socketChannel.close();
                            } else if (bytesRead > 0) {
                                buffer.flip();

                                if(attachment instanceof Integer){
                                    //协议和认证确定，然后连接

                                    int i = (Integer) attachment;
                                    if(i == 1){
                                        //第一步 协商连接
                                        negotiatedConnection(socketChannel,buffer);
                                        key.attach(2);
                                    }else if(i == 2){
                                        //第二步 连接客户端发送的ip和端口
                                        connectServer(selector,socketChannel,buffer);
                                        key.attach(3);
                                    }

                                }else {
                                    //一个channel写入另一个channel
                                    SocketChannel clientChannel =  (SocketChannel) attachment;
                                    if(clientChannel.isConnected()){
                                        clientChannel.write(buffer);
                                    }

                                }

                            }
                        }catch (Exception e){

                            e.printStackTrace();
                            socketChannel.close();

                        }

                    }
                }catch (Exception e){
                    e.printStackTrace();
                }
            }
            selector.selectedKeys().clear();
        }
    }
    private static void negotiatedConnection(SocketChannel socketChannel,ByteBuffer buffer) throws IOException {
        //SOCKS版本
        byte version = buffer.get();
        //认证方法数量
        byte nmethods = buffer.get();
        //认证方法
        boolean noAuth = false;
        byte[] methods = new byte[nmethods];
        for (byte b = 0; b < nmethods; b++) {
            methods[b] = buffer.get();
            //无需验证
            if(methods[b] == 0){
                noAuth = true;
            }
        }

        //协议是socks5，支持无需要验证
        if(version == 5 && noAuth){
            //回复客户端,版本5，无需验证
            byte[] reply = {5,0};
            socketChannel.write(ByteBuffer.wrap(reply));
        }else{
            //不支持关闭连接
            socketChannel.close();
        }
    }
    private static void connectServer(Selector selector,SocketChannel socketChannel, ByteBuffer buffer) throws IOException {
        //SOCKS版本
        byte version = buffer.get();
        //CMD：命令，1 表示连接，2 表示绑定。
        byte cmd = buffer.get();
        //保留字段，必须为 0x00。
        byte rsv = buffer.get();
        //地址类型，1 表示 IPv4 地址，3 表示域名，4 表示 IPv6 地址
        byte atyp = buffer.get();
        String host = "";
        if(atyp == 1){
            //ipv4
            byte[] bt = new byte[4];
            buffer.get(bt);
            host = InetAddress.getByAddress(bt).getHostAddress();

        }else if(atyp == 3){
            //域名
            byte len = buffer.get();
            byte[] bt = new byte[len];
            buffer.get(bt);
            host = new String(bt);

        }else if(atyp == 4){
            //ipv6
            byte[] bt = new byte[16];
            buffer.get(bt);
            host = InetAddress.getByAddress(bt).getHostAddress();
        }

       int port = Short.toUnsignedInt(buffer.getShort());

        System.out.println("host:"+host+",port:"+port);

        InetSocketAddress address = new InetSocketAddress(host, port);
        SocketChannel clientChannel = SocketChannel.open();
        clientChannel.configureBlocking(false);
        clientChannel.connect(address);

        byte[] replyData = Arrays.copyOf(buffer.array(), buffer.position());
        //应答成功
        replyData[1] = 0;
        Object[] array = new Object[2];
        array[0] = socketChannel;
        array[1] = replyData;
        clientChannel.register(selector,SelectionKey.OP_CONNECT,array);

    }


}